package com.ddys.triemap;

import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * Author <a href="mailto:xiaowenke@zto.cn">xwk</a>
 * Date: 2016-08-25
 * Time: 10:17
 * To change this template use File | Settings | File Templates.
 */
public class TrieMap<T> {
    private TrieNode<T> root;

    public TrieMap() {
        root = new TrieNode<>();
    }

    public void insert(T value, String... words) {
        if (words.length <= 0) {
            return;
        }
        for (String word : words) {
            if (word == null) {
                continue;
            }
            TrieNode<T> current = root;
            for (int i = 0; i < word.length(); i++) {
                Character c = word.charAt(i);
                if (!current.isChild(c)) {
                    current.putChild(c);
                }
                current = current.getChild(c);
                if (i == word.length() - 1) {
                    current.addValue(value);
                }

            }
        }
    }

    public List<T> search(String word, Criteria criteria) {
        List<T> list = new ArrayList<T>();
        TrieNode lastNode = findLastNode(word);
        loadChild(list, lastNode, criteria);
        return list;
    }
    public long countSearch(String word, Criteria criteria){
        if(criteria.getFilter()==null){
            TrieNode current = findLastNode(word);
            if(current!=null){
                return current.getSize();
            }else{
                return 0;
            }
        }else{
            Criteria c = new Criteria(criteria.getFilter());
            List<T> result = search(word,c);
            return result.size();
        }
    }
    private TrieNode findLastNode(String word){
        if (!(word == null || word.length() <= 0)) {
            TrieNode current = root;
            for (int i = 0; i < word.length(); i++) {
                if (current.getChild(word.charAt(i)) == null) {
                    return null;
                }
                current = current.getChild(word.charAt(i));
            }
            return current;
        }else{
            return null;
        }
    }

    private void loadChild(List<T> list, TrieNode current, Criteria criteria) {
        if(current==null){
            return;
        }
        Set<T> values = current.getValues();
        if (values != null) {
            for (T value : values) {
                Criteria.STATUS status = criteria.next(value);
                if (status == Criteria.STATUS.DONE) {
                    return;
                } else if (status == Criteria.STATUS.HIT) {
                    list.add(value);
                }
            }

        }
        Map<Character, TrieNode<T>> childs = current.getChilds();
        if (childs != null && childs.size() > 0) {
            for (Map.Entry<Character, TrieNode<T>> e : childs.entrySet()) {
                current = e.getValue();
                loadChild(list, current, criteria);
            }
        }
    }

    class TrieNode<T> {
        private long size = 0L;
        private Map<Character, TrieNode<T>> nodeMap;
        private Set<T> values;

        public long getSize() {
            return size;
        }

        public void addValue(T value) {
            if (values == null) {
                values = new HashSet<T>();
            }
            values.add(value);
            this.size++;
        }

        public boolean isChild(Character c) {
            if (this.nodeMap == null) {
                return false;
            }
            return this.nodeMap.containsKey(c);
        }

        public void putChild(Character c) {
            if (this.nodeMap == null) {
                this.nodeMap = new HashMap<>();
            }
            this.nodeMap.put(c, new TrieNode<T>());
            this.size++;
        }

        public TrieNode<T> getChild(Character c) {
            return this.nodeMap.get(c);
        }

        public Map<Character, TrieNode<T>> getChilds() {
            return this.nodeMap;
        }

        public Set<T> getValues() {
            return values;
        }

    }


}
